import os
import cv2
import json
import random
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from data_utils import transform as transform_v1
from image_augment import transform as transform_v2, mixup

class PascalVOCDatasetV1(keras.utils.Sequence):
    """
        进入的box是x1y1x2y2的真实坐标
        返回的box是x1y1x2y2的比例格式
    """
    def __init__(self, batch_size, data_folder, split, output_shape=(300,300), transform=None, keep_difficult=False, shuffle=False):
        """
            input_shape: (w, h)
        """
        self.split = split
        assert self.split in {'TRAIN', 'TEST'} #! SET
        self.data_folder = data_folder
        self.keep_difficult = keep_difficult
        self.output_shape = output_shape
        self.transform = transform

        with open(os.path.join(data_folder, self.split + "_images.json")) as f:
            self.raw_images = json.load(f)
            self.images = self.raw_images

        with open(os.path.join(data_folder, self.split + '_objects.json')) as f:
            self.raw_objects = json.load(f)
            self.objects = self.raw_objects
        
        assert len(self.images) == len(self.objects)

        self.batch_size = batch_size

        self.shuffle = shuffle

        if self.shuffle:
            self.indices = list(range(len(self.raw_images)))
            self.random()

    def __len__(self):

        return int(np.ceil( len(self.images) / float(self.batch_size) ))
    
    def __getitem__(self, i):

        batch_image_path  = self.images [i * self.batch_size: (i+1) * self.batch_size]
        batch_objects     = self.objects[i * self.batch_size: (i+1) * self.batch_size]

        batch_images = list()
        batch_boxes  = list()
        batch_labels = list()
        batch_difficulties = list()

        for image_path, objects in zip(batch_image_path, batch_objects):
            image = cv2.imread(image_path)
            image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)

            boxes  = np.array(objects['boxes'],dtype=np.float32)
            labels = np.array(objects['labels'],dtype=np.int64)
            difficulties = np.array(objects['difficulties'],dtype=np.uint8)

            if not self.keep_difficult:
                boxes  = boxes[np.where(1-difficulties)]
                labels = labels[np.where(1-difficulties)]
                difficulties = difficulties[np.where(1-difficulties)]

            split = self.split if self.transform is None else ("TRAIN" if self.transform else "TEST")

            image, boxes, labels, difficulties = transform_v1(image, boxes, labels, difficulties, split=split, wh=self.output_shape, bgr=False)

            batch_images.append(image)
            batch_boxes.append(boxes)
            batch_labels.append(labels)
            batch_difficulties.append(difficulties)

        batch_images = np.stack(batch_images)

        return batch_images, batch_boxes, batch_labels, batch_difficulties
    
    def random(self):
        random.shuffle(self.indices)
        self.images  = [self.raw_images [ind] for ind in self.indices]
        self.objects = [self.raw_objects[ind] for ind in self.indices]

    def on_epoch_end(self):
        if self.shuffle: self.random()

class PascalVOCDatasetV2(PascalVOCDatasetV1):
    def __init__(self, batch_size, data_folder, split, output_shape=(300,300), transform=None, keep_difficult=False, shuffle=False, mix_batch_size=16):
        super(PascalVOCDatasetV2, self).__init__(batch_size, data_folder, split, output_shape=output_shape, transform=transform, keep_difficult=keep_difficult, shuffle=shuffle)
        self.mix_batch_size = mix_batch_size

    def __getitem__(self, i):

        batch_image_path  = self.images [i * self.batch_size: (i+1) * self.batch_size]
        batch_objects     = self.objects[i * self.batch_size: (i+1) * self.batch_size]

        batch_images = list()
        batch_boxes  = list()
        batch_labels = list()
        batch_difficulties = list()

        MIX_BATCH_SIZE = self.mix_batch_size

        split = self.split if self.transform is None else ("TRAIN" if self.transform else "TEST")

        if self.batch_size >= MIX_BATCH_SIZE: #! batch_size大于等于MIX_BATCH_SIZE时, 默认使用mixup方法
            image_list  = list()
            boxes_list  = list()
            labels_list = list()
            difficulties_list = list()

        #! 数据处理和增强
        for image_path, objects in zip(batch_image_path, batch_objects):
            image = cv2.imread(image_path)
            image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)

            boxes  = np.array(objects['boxes'],dtype=np.float32)
            labels = np.array(objects['labels'],dtype=np.int64)
            difficulties = np.array(objects['difficulties'],dtype=np.uint8)

            if not self.keep_difficult:
                boxes  = boxes[np.where(1-difficulties)]
                labels = labels[np.where(1-difficulties)]
                difficulties = difficulties[np.where(1-difficulties)]

            if self.batch_size >= MIX_BATCH_SIZE: #! batch_size大于等于MIX_BATCH_SIZE时, 默认使用mixup方法
                image_list.append(image)
                boxes_list.append(boxes)
                labels_list.append(labels)
                difficulties_list.append(difficulties)

            image, boxes, labels, difficulties = transform_v2(image, boxes, labels, difficulties, split=split, new_wh=self.output_shape, bgr=False)

            batch_images.append(image)
            batch_boxes.append(boxes)
            batch_labels.append(labels)
            batch_difficulties.append(difficulties)

        #! 数据拼接
        if split == "TRAIN":

            num = min(self.batch_size // MIX_BATCH_SIZE, len(batch_images) // MIX_BATCH_SIZE) #! 为避免数据索引溢出

            i = 0

            while i < num:
                sub_indice = random.randint(i * MIX_BATCH_SIZE, min((i+1) * MIX_BATCH_SIZE - 1, i * MIX_BATCH_SIZE + len(batch_images) -1))
                if random.random() < 0.5: #! 缩放拼接
                    image, boxes, labels, difficulties = \
                        mixup(image_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            boxes_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            labels_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            difficulties_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE],
                            self.output_shape)
                else: #! 剪裁拼接
                    image, boxes, labels, difficulties = \
                        mixup(image_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            boxes_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            labels_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE], 
                            difficulties_list[i*MIX_BATCH_SIZE: (i+1)*MIX_BATCH_SIZE],
                            self.output_shape,
                            crop=True)
                
                if not boxes.any(): continue
                else: i+=1

                batch_images[sub_indice] = image
                batch_boxes[sub_indice] = boxes
                batch_labels[sub_indice] = labels
                batch_difficulties[sub_indice] = difficulties

        batch_images = np.stack(batch_images)

        return batch_images, batch_boxes, batch_labels, batch_difficulties